# Android environment
import time
import traceback
import json
import os
import docker
import selenium
import random
import base64
import time
from appium import webdriver
from enum import Enum
from copy import deepcopy
from typing import Dict, Optional, List, Tuple
from docker.models.containers import Container
from docker.errors import NotFound
from abc import ABCMeta, abstractmethod
from appium.webdriver.webdriver import WebDriver
from selenium.common.exceptions import WebDriverException
from selenium.common.exceptions import NoSuchElementException
from appium.webdriver.webelement import WebElement
from androidsandbox.utils import logger, Utils, CFG
from mispconn.apk_parse3.apk import APK
import requests.exceptions
import subprocess
import os
import shlex
from threading import Timer
from selenium.common.exceptions import NoSuchElementException


class DeviceError(Exception):
    def __init__(self, msg):
        super().__init__()
        self.msg = msg

    def __str__(self):
        return repr(self.msg)


class DeviceType(Enum):
    """Types of Android devices. """
    DockerEmulator = 1 # budtmo/docker-android-x86-11.0 (with Appium built in)
    HostEmulator = 2  # Emulators on host
    Real = 3  # Real devices

class Device(metaclass=ABCMeta):
    """Android device with appium and adb controller"""
    def __init__(self, appium_ip: str, host_dir: str, udid: str,
                 appium_port: int = 4723):
        """
        :param appium_ip: IP address of Appium server.
        :param host_dir: directory on the machine that runs adb server.
        :param udid: serial number in adb.
        :param appium_port: port of Appium server, default to 4723.
        """
        self.udid: str = udid  # serial number
        self.appium_ip: str = appium_ip
        self.appium_port: int = appium_port
        self.host_dir: str = host_dir
        self.webdriver: Optional[WebDriver] = None  # WebDeriver of Appium
        # Files such as pcap are stored in data_dir_on_android.
        self.data_dir_on_android: str = 'data/myuitest'

    def build_webdriver(self, sample_file_name: str, mode: str = 'basic'):
        """
        Connect to Appium server.
        :param sample_file_name:
        :param mode: basic or uitest
        :return:
        """
        desired_capabilities: Dict = deepcopy(CFG['desired_capabilities'])
        # 增加超时设置
        desired_capabilities['uiautomator2ServerLaunchTimeout'] = 60000  # 例如设置为60秒
        if mode == 'uitest':
            sample_path = os.path.join(self.host_dir, sample_file_name)
            desired_capabilities['app'] = sample_path
        elif mode == 'basic':
            pass
        logger.info(f'building WebDriver ({mode})...')

        def webdriver_built_func():
            try:
                self.webdriver = webdriver.Remote(
                    f'http://{self.appium_ip}:{self.appium_port}/wd/hub',
                    desired_capabilities, keep_alive=False)
                logger.info(f'built WebDriver ({mode}) successfully.')
                return True
            except WebDriverException as e:
                logger.error(f'Error building WebDriver ({mode}): {e}')
                logger.debug(traceback.format_exc())
                return False
        if not Utils.wait_until(webdriver_built_func, 1):
            raise DeviceError(f'failed to built WebDriver ({mode}).')

    def ensure_availability(self) -> bool:
        """
            Ensure that the device is available by actions such as restarting.
            Return True if the device is available, False otherwise.
        """
        count = 0
        while not self.is_available():
            # Restart the docker container a maximum of limited times.
            if count >= CFG['max_times_device_restart']:
                logger.critical(f'device {self.udid} is not available.')
                return False
            self.restart()
            count += 1
        else:
            return True

    @abstractmethod
    def is_available(self) -> bool:
        pass

    def renew(self):
        """Destroy self if dead and create a new one"""
        logger.info('renewing device ...')
        self.destroy()
        device = self.create()
        if device and device.ensure_availability():
            return device
        else:
            return None

    @abstractmethod
    def destroy(self):
        """Destroy self"""
        pass

    @abstractmethod
    def restart(self):
        pass

    @staticmethod
    @abstractmethod
    def create():
        """Create a new instance"""
        pass

    @abstractmethod
    def disable_selinux(self):
        pass

    @abstractmethod
    def strace(self, packagename, sample_md5):
        """Trace system calls"""
        pass

    @abstractmethod
    def run_adbd_as_root(self):
        pass

    def test_app(self, sample_md5: str, package_name: str):
        """
        :param package_name: package name of the app.
        :param sample_md5: md5 of the sample file.
        :return:
        """
        sample_file_name = f'{sample_md5}.apk'
        self.run_adbd_as_root()
        try:
            self.build_webdriver(sample_file_name, mode='uitest')
        except DeviceError:
            # Though failed to build webdriver with app, the app may already
            # been installed.
            self.build_webdriver(sample_file_name, mode='basic')
            if package_name not in self.list_3rd_packages():
                raise DeviceError('app is not installed.')
            
        self.disable_selinux()
        self.stop_all_tcpdump()
        self.remove_sample_data()
        self.remove_3rd_packages(package_name)
        self.make_sample_data_dir()
        self.start_tcpdump(sample_md5)
        self.strace(package_name, sample_md5)
        self.start_ui_test(sample_md5, package_name)
        self.stop_all_tcpdump()
        self.pull_pcap(sample_md5)
        self.pull_strace(sample_md5)
        self.remove_3rd_packages()

    def list_3rd_packages(self):
        result = []
        try:
            outs = self.webdriver.execute_script(
                'mobile:shell', {'command': 'pm',
                                 'args': ['list', 'packages', '-3'],
                                 'timeout': 2 * 60 * 1000})
            if type(outs) is str:
                result = [pkg.split(':')[1] for pkg in outs.strip().split('\n')
                          if 'appium' not in pkg]
        except WebDriverException:
            raise DeviceError('failed to list third party packages')
        return result

    def remove_3rd_packages(self, package_name: Optional[str] = None):
        """Remove all third party packages except
        package_name before starting a new test."""
        third_packages = self.list_3rd_packages()
        for pkg in third_packages:
            # For android 5.0.1, some system packets are listed as 3rd ones
            if CFG['device_types']['docker']['image'] == \
                    "budtmo/docker-android-x86-5.0.1" and pkg in \
                    ['com.example.android.apis\r',
                     'com.example.android.softkeyboard\r',
                     'com.android.smoketest.tests\r', 'com.android.smoketest\r',
                     'com.android.widgetpreview\r',
                     'com.example.android.livecubes\r',
                     'com.android.gesture.builder\r']:
                continue
            if package_name is not None and package_name in pkg:
                continue
            try:
                self.webdriver.execute_script(
                    'mobile:shell', {'command': 'pm',
                                     'args': ['uninstall',
                                              '--user', '0', pkg],
                                     'timeout': 60 * 1000})
                logger.info(f'uninstalled third party package '
                            f'{pkg} successfully.')
            except WebDriverException:
                raise DeviceError(f'failed to uninstall '
                                  f'third party package: {pkg}')

    def install_app(self, sample_file_name: str):
        sample_path = os.path.join(self.host_dir, sample_file_name)
        try:
            self.webdriver.install_app(sample_path)
            logger.info(f'installed {sample_file_name} successfully.')
        except:
            raise DeviceError(f'failed to install {sample_file_name}')


    def start_ui_test(self, sample_md5: str, package_name: str):
        """
        Automatically click items and take screenshots.
        :param package_name: package name of the app.
        :param sample_md5: md5 of the sample file
        :return:
        """
        driver = self.webdriver
        driver.update_settings({'ignoreUnimportantViews': True})

        # Get the context
        context = {'sample_md5': sample_md5}
        device_info = driver.execute_script('mobile:deviceInfo')
        for key in ("apiVersion", "platformVersion", "brand", "model",
                    "realDisplaySize", "timeZone"):
            context[key] = device_info[key]
        context['bars'] = driver.get_system_bars()
        context['screenshots'] = []
        context['pages'] = []
        depth_cursor = 0     # current UI depth
        page_cursor = 0         # current page id
        screenshot_id_generator = Utils.gen_id()
        page_id_generator = Utils.gen_id()

        if isinstance(self, DockerDevice):
            screenshot_dir = self.super_host_dir

        def click(element: WebElement) -> bool:
            """
            Click and update cursors.
            :param element: which is clicked.
            :return: whether the page after clicking is new.
            """
            nonlocal depth_cursor
            element.click()
            depth_cursor += 1
            time.sleep(CFG['interval_between_clicks'])
            _, is_new_page = update_page_cursor()
            return is_new_page

        def safe_get_page_source():
            """
            安全地获取页面源码，附带异常处理
            """
            try:
                return driver.page_source
            except WebDriverException as e:
                logger.error(f'获取页面源码时发生错误: {e}')
                logger.debug(traceback.format_exc())
                return ''
        
        def update_page_cursor() -> Tuple[int, bool]:
            """
            Update page cursor based on current page.
            :return: (page id, is a new page)
            """
            nonlocal driver, page_cursor, page_id_generator, depth_cursor
            page = safe_get_page_source()
            for p in context['pages']:
                if p['source'] == page:
                    # The page is already there.
                    page_cursor = p['id']
                    return p['id'], False
            # A new page.
            new_page_id = next(page_id_generator)
            context['pages'].append({
                "id": new_page_id,
                "source": page,
                "depth": depth_cursor
            })
            page_cursor = new_page_id
            return new_page_id, True

        def go_back():
            """Go back after clicking."""
            nonlocal driver, depth_cursor
            driver.back()
            depth_cursor -= 1
            update_page_cursor()

        def get_element_attrs(element: WebElement) -> Dict:
            return {
                'classname': element.get_attribute('classname'),
                'text': element.get_attribute('text'),
                'resource_id': element.get_attribute('resource-id')
            }

        def take_screenshot():
            # Take a screenshot without clicking any element.
            nonlocal screenshot_id_generator, depth_cursor, screenshot_dir
            nonlocal package_name
            launch_with_overcoming_older_app()
            update_page_cursor()
            filename = f'{depth_cursor}-{next(screenshot_id_generator)}.png'
            assert self.webdriver.current_package == package_name
            driver.get_screenshot_as_file(f'{screenshot_dir}/{filename}')
            screenshot_context = {
                "filename": filename,
                "device_time": driver.get_device_time(),
                "current_package": driver.current_package,
                "current_activity": driver.current_activity,
                "entrance":
                    {
                        "page_id": None,
                        "element": None
                    },
                "page_id": page_cursor
            }
            context['screenshots'].append(screenshot_context)

        def click_and_take_screenshot(elmt: WebElement):
            """
            Click element, save the page, take a screenshot and then go back.
            Page cursor and depth cursor are both updated accordingly.
            :param elmt: the element which is clicked.
            """
            nonlocal depth_cursor, page_cursor, screenshot_id_generator, screenshot_dir
            nonlocal package_name

            try:
                # 确保应用已启动
                launch_with_overcoming_older_app()
                
                # 获取点击前元素的属性
                element_attrs = get_element_attrs(elmt)
                last_page_id = page_cursor

                # 执行点击操作，并检查是否导航到了新页面
                is_new_page = click(elmt)
                if not is_new_page:
                    go_back()
                    return
                # 拍摄截图
                filename = f'{depth_cursor}-{next(screenshot_id_generator)}.png'
                assert self.webdriver.current_package == package_name
                driver.get_screenshot_as_file(f'{screenshot_dir}/{filename}')
                # 记录截图上下文
                screenshot_context = {
                    "filename": filename,
                    "device_time": driver.get_device_time(),
                    "current_package": driver.current_package,
                    "current_activity": driver.current_activity,
                    "entrance": {
                        "page_id": last_page_id,
                        "element": element_attrs
                    },
                    "page_id": page_cursor
                }
                context['screenshots'].append(screenshot_context)
                # 返回之前的页面
                go_back()
            except NoSuchElementException as e:
                logger.error(f'未找到元素: {e}')
                logger.debug(traceback.format_exc())
            except WebDriverException as e:
                logger.error(f'点击元素并截图过程中发生 WebDriver 错误: {e}')
                logger.debug(traceback.format_exc())
            except Exception as e:
                logger.error(f'在 click_and_take_screenshot 函数中发生未知错误: {e}')
                logger.debug(traceback.format_exc())

        def save_context():
            nonlocal context
            host_dir = self.host_dir
            if isinstance(self, DockerDevice):
                host_dir = self.super_host_dir
            with open(f'{host_dir}/context.json', 'w') as f:
                json.dump(context, f)
            logger.info('saved context.json successfully.')

        def launch_app():
            nonlocal package_name
            cmd = f'cmd package resolve-activity --brief {package_name} | ' \
                  f'tail -n 1'
            main_activity = self.webdriver.execute_script(
                'mobile:shell', {'command': cmd, 'timeout': 20 * 1000}).strip()
            cmd = f'am start {main_activity}'
            logger.debug(f'launching f{main_activity}')
            self.webdriver.execute_script(
                'mobile:shell', {'command': cmd, 'timeout': 20 * 1000})
            time.sleep(CFG["interval_between_clicks"])

        def launch_with_overcoming_older_app():
            """New android will prompt on apps built for older version."""
            while "This app was built for an older version of Android" \
                    in self.webdriver.page_source:
                self.webdriver.find_element_by_android_uiautomator(
                    'new UiSelector().textContains("OK").clickable(true)').\
                    click()
                if self.webdriver.current_package != package_name:
                    launch_app()
                    # self.webdriver.launch_app()
            if self.webdriver.current_package != package_name:
                launch_app()
                # self.webdriver.launch_app()
            time.sleep(CFG['interval_between_clicks'])
            assert self.webdriver.current_package == package_name
        

        # running UI test
        logger.info('starting UI test...')
        screenshot_count = 0

        # 检测是否有弹窗。
        try:
            self.webdriver.find_element_by_android_uiautomator('new UiSelector().text("Wait").clickable(true)').click()
            logger.info('clear windos : Wait')
        except :
            logger.error('No find windows')
        
        # boot screen
        time.sleep(5)  # ignore loading page
        try:
            take_screenshot()
            screenshot_count += 1
            logger.info('captured boot screen successfully.')
        except:
            logger.error(f'failed to capture the boot screen.')
            logger.debug(traceback.format_exc())
            raise DeviceError(f'failed to start the app.')
        # agreement screen
        try:
            for word in CFG['agreement_words']:
                selector = f'new UiSelector().textContains("{word}")' \
                           f'.clickable(true)'
                try:
                    elm = driver.find_element_by_android_uiautomator(selector)
                    click_and_take_screenshot(elm)
                    screenshot_count += 1
                    logger.info('captured agreement screen successfully.')
                    break
                except NoSuchElementException:
                    pass
        except:
            logger.error(f'failed to capture the agreement screen.')
            logger.debug(traceback.format_exc())
        # Randomly click an element and go back. Note that the back action may
        # not navigate to the exactly previous page. If so, the depth of page
        # in context is not accurate.
        click_count = 0
        while click_count < CFG['amount_clicks']:
            try:
                clickable_elements = []
                for element in driver.find_elements_by_android_uiautomator(
                    f'new UiSelector().packageName("{package_name}")'):
                        if element.get_attribute("clickable") == 'true':
                            clickable_elements.append(element)
                if not clickable_elements:
                    break
                chosen_element = random.choice(clickable_elements)
                click_and_take_screenshot(chosen_element)
                screenshot_count += 1
                logger.info(f'captured screen for click '
                            f'{click_count} sscuccessfully.')
            except Exception as e:
                logger.error(f'failed to capture the screen '
                             f'for click {click_count}.')
                logger.debug(traceback.format_exc())
                if 'socket hang up' in str(e):
                    break
            finally:
                click_count += 1
        if screenshot_count > 0:
            save_context()
            logger.info(f'finished UI test with {screenshot_count} '
                        f'screenshots captured.')
        else:
            raise DeviceError('failed to capture any screenshot.')

    def make_sample_data_dir(self):
        try:
            self.webdriver.execute_script('mobile:shell',
                                          {'command': 'mkdir',
                                           'args': [self.data_dir_on_android],
                                           'timeout': 2 * 60 * 1000})
            logger.info(f'made {self.data_dir_on_android} '
                        f'on android successfully.')
        except WebDriverException as e:
            if 'File exists' in str(e):
                pass
            else:
                raise DeviceError('failed to make sample data directory.')

    def start_tcpdump(self, sample_md5: str):
        """Start tcpdump on android."""
        try:
            self.webdriver.execute_script(
                'mobile:shell',
                {'command': 'tcpdump',
                 'args': ['tcp or udp and not udp port 5353',
                          '-s0', '-n', '--packet-buffered', '--immediate-mode',
                          '-w', f'{self.data_dir_on_android}/{sample_md5}.pcap',
                          '&'],
                 'timeout': 5 * 1000})
            #  are not available on
            # all sdk versions
        except WebDriverException:
            pass

        if not self.is_tcpdump_running():
            raise DeviceError('failed to start tcpdump.')
        else:
            logger.info('started tcpdump successfully.')

    def is_tcpdump_running(self):
        try:
            if CFG['device_types']['docker']['image'] == \
                    "budtmo/docker-android-x86-5.0.1":
                ps = self.webdriver.execute_script(
                    'mobile:shell', {'command': 'ps | grep "tcpdump"'})
                if type(ps) and 'tcpdump' in ps:
                    return True
                else:
                    return False
            else:
                ps = self.webdriver.execute_script(
                    'mobile:shell', {'command': 'ps -elf | grep "tcpdump"'})
                if type(ps) is str and len(ps.strip().split('\n')) >= 3:
                    return True
                else:
                    return False
        except:
            return False

    def is_data_dir_existing(self):
        try:
            self.webdriver.execute_script(
                'mobile:shell', {'command': 'ls',
                                 'args': [self.data_dir_on_android]})
            return True
        except WebDriverException:
            return False

    def remove_sample_data(self):
        def sample_data_removed_func():
            if not self.is_data_dir_existing():
                return True
            try:
                self.webdriver.execute_script(
                    'mobile:shell', {'command': 'rm',
                                     'args': ['-rf',
                                              f'{self.data_dir_on_android}/*']})
                logger.info(f'cleared {self.data_dir_on_android} '
                            f'on android successfully')
                return True
            except WebDriverException:
                return False

        if not Utils.wait_until(sample_data_removed_func, 5):
            raise DeviceError('failed to remove sample data.')

    def stop_all_tcpdump(self):
        """Stop all tcpdump processes on android."""
        def all_tcpdump_are_killed_func():
            if not self.is_tcpdump_running():
                return True
            try:
                self.webdriver.execute_script(
                    'mobile:shell', {'command': 'killall tcpdump'})
            except:
                pass
            return False

        if not Utils.wait_until(all_tcpdump_are_killed_func, 20):
            raise DeviceError('failed to stop tcpdump.')
        logger.info('killed all tcpdump successfully.')

    @abstractmethod
    def pull_pcap(self, sample_md5: str):
        pass

    @abstractmethod
    def pull_strace(self, sample_md5: str):
        pass


class DockerDevice(Device):
    def __init__(self, container_name: str, appium_ip: str,
                 host_dir: str, super_host_dir: str, udid: str,
                 appium_port: int = 4723):
        """
        :param container_name:
            only make sense for DockerEmulator.
        :param appium_ip:
            IP address of Appium server.
        :param host_dir:
            directory on the machine that runs adb server, such as a docker
            container.
        :param super_host_dir:
            directory on the machine that runs this program. host_dir and
            super_host_dir are identical for read and emulator, but different
            for docker.
        :param udid:
            serial number in adb.
        :param appium_port:
            port of Appium server, default to 4723.
        """
        super().__init__(appium_ip, host_dir, udid, appium_port)
        self.super_host_dir = super_host_dir
        self.container_name: str = container_name

    def is_available(self):
        """Whether the device is available now."""
        android_available = False
        appium_available = False
        # Check Android device's availability.
        cmd = f'docker exec {self.container_name} ' \
              f'"/root/platform-tools/adb" "devices"'
        if self.udid in Utils.run_cmd(cmd)[0]:
            android_available = True
        else:
            logger.error('failed to run adb.')
        # Check Appium server's availability.
        if selenium.webdriver.common.utils.is_connectable(self.appium_port,
                                                          host=self.appium_ip):
            appium_available = True
        else:
            logger.error('failed to connect to appium.')
        if android_available and appium_available:
            return True
        else:
            logger.warning(f'{self.container_name}:{self.udid} '
                           f'is not available.')
            return False

    @classmethod
    def create(cls) -> Optional[Device]:
        device = None
        suffix = Utils.get_random_alphanumeric_string(8)
        super_host_dir = f'/tmp/{suffix}'
        host_dir = '/root/tmp'
        os.makedirs(super_host_dir, exist_ok=True)
        container_name = f'android-container-{suffix}'
        logger.info(f'creating {container_name} ...')
        docker_image = CFG['device_types']['docker']['image']
        cmd = f'docker run --privileged -d ' \
              f'-e DEVICE="Nexus 5" ' \
              f'-e APPIUM=true -e RELAXED_SECURITY=true ' \
              f'-v {super_host_dir}:{host_dir} ' \
              f'--name {container_name} {docker_image}'
        Utils.run_cmd(cmd)
        # if not outs or len(outs.strip()) != 64:
        #     # A succeeded run will return a sha256 string
        docker_client = docker.from_env()

        def container_created_func():
            try:
                docker_client.containers.get(container_name)
                return True
            except NotFound:
                return False

        Utils.wait_until(container_created_func, 300)
        container: Container = docker_client.containers.get(container_name)
        if container is None:
            logger.error(f'failed to create {container_name}.')
        else:
            device = DockerDevice.extract_device_from_container(container)
        return device

    def disable_selinux(self):
        try:
            enforce = self.webdriver.execute_script('mobile:shell',
                                                    {'command': 'cat',
                                                     'args': '/sys/fs/selinux/enforce',
                                                     'timeout': 2 * 60 * 1000})
            if enforce == '1':
                self.webdriver.execute_script('mobile:shell',
                                              {'command': 'setenforce',
                                               'args': '0',
                                               'timeout': 2 * 60 * 1000})
                enforce = self.webdriver.execute_script('mobile:shell',
                                                        {'command': 'cat',
                                                         'args': '/sys/fs/selinux/enforce',
                                                         'timeout': 2 * 60 * 1000})
            assert enforce == '0'
            logger.info(f'SELinux disabled.')
        except Exception as e:
            raise DeviceError(f'failed to disable SELinux {e}.')


    def wait_for_process(self,container_name,package_name, timeout=10):
        start_time = time.time()
        while time.time() - start_time < timeout:
            result = subprocess.run(f'docker exec -t {container_name} "adb" "shell"  "pidof {package_name}"', shell=True, capture_output=True, text=True)
            if result.stdout:
                print(result.stdout.strip())
                return result.stdout.strip().split()[0]
            time.sleep(0.1)
        return None
    
    def strace(self, package_name, sample_md5):
        """Trace system calls. The trace collection is independent cuz appium test will not produce any syscall trace"""
        try:
            print('开始strace')
            sample_file_name = f'{sample_md5}.apk'
            sample_path = os.path.join(self.super_host_dir, sample_file_name)
            main_activity = APK(sample_path).get_main_activity()
            
            # 杀死所有程序
            cmd = f'docker exec -t {self.container_name} "adb" "shell" "killall {package_name}"'
            Utils.run_cmd(cmd)

            # 启动进程。
            # cmd = f'am start -n {package_name}/{main_activity} && sleep 1 && set `pidof {package_name}` && strace -t -tt -y -yy -s 1024 -f -p $1 -o /{self.data_dir_on_android}/{sample_md5[:16]}.strace'
            # 原来的命令
            # cmd = f'docker exec -t {self.container_name} "adb" "shell" "am start -n {package_name}/{main_activity} && sleep 0.2 && set `pidof {package_name}` && strace -t -tt -y -yy -s 128 -f -p $1 -o /{self.data_dir_on_android}/{sample_md5[:16]}.strace"'
            cmd = f'docker exec -t {self.container_name} "adb" "shell" "am start -n {package_name}/{main_activity}"'
            Utils.run_cmd(cmd)
            # 获取进程ID
            pid = self.wait_for_process(self.container_name,package_name)
            print(pid)
            if pid:
                cmd = f'docker exec -t {self.container_name}  "adb" "shell" "strace -t -tt -y -yy -s 128 -f -p {pid} -o /{self.data_dir_on_android}/{sample_md5[:16]}.strace"'
                Utils.run_cmd(cmd)
            else:
                print("未能找到主进程的PID")
            
            time.sleep(180)
            cmd = f'docker exec -t {self.container_name} "adb" "shell" "killall strace"'
            Utils.run_cmd(cmd)
            print('结束strace')
            # time.sleep(20)
            # # 构造用于检查文件存在性的命令
            # cmd_check_file = f'docker exec -t {self.container_name} "adb" "shell" "test -e /{self.data_dir_on_android}/{sample_md5[:16]}.strace" && echo Exists || echo Does not exist"'
            # # 执行命令
            # exists = Utils.run_cmd(cmd_check_file)
            # print(exists)


            # cmd = f'docker exec -t {self.container_name} "adb" "shell" "tar -C /{self.data_dir_on_android}/{sample_md5[:16]}.strace -czf /{self.data_dir_on_android}/{sample_md5}.strace.tar.gz ."'
            # Utils.run_cmd(cmd)
            # self.webdriver.execute_script('mobile:shell',
            #                             {'command': 'setprop',
            #                              'args': [f'wrap.{package_name}', f'"logwrapper strace -t -tt -y -yy -s 1024 -f -o /{self.data_dir_on_android}/{sample_md5[:16]}.strace"'],
            #                              'timeout': 2 * 60 * 1000})
            # self.webdriver.execute_script('mobile:shell',
            #                               {'command': 'monkey',
            #                                'args': ['-p', package_name, '10'],
            #                                'timeout': 2 * 60 * 1000})
            
            logger.info(f'strace collected.')
            print("strace 收集完咯！")
        except Exception as e:
            raise DeviceError(f'failed to collect strace {e}.')

    def restart(self):
        docker_client = docker.from_env()
        container: Container = docker_client.containers.get(self.container_name)
        logger.info(f'restarting {self.container_name} ...')
        Utils.run_cmd(f'docker restart {self.container_name}')

        def healthy_func():
            try:
                container.reload()
                return container.attrs['State']['Health']['Status'] == 'healthy'
            except:
                return False

        if Utils.wait_until(healthy_func, timeout=900):
            pass
        else:
            msg = f'failed to restart {self.container_name}.'
            raise DeviceError(msg)

    def destroy(self):
        try:
            cmd = f'docker stop {self.container_name}'
            Utils.run_cmd(cmd, timeout=600)
            cmd = f'docker rm {self.container_name}'
            Utils.run_cmd(cmd, timeout=600)
            logger.info(f'destroyed {self.container_name}.')
        except:
            raise DeviceError(f'failed to destroy {self.container_name}')

    @classmethod
    def extract_device_from_container(cls, container: Container):
        return DockerDevice(
            container_name=container.name,
            host_dir=container.attrs['Mounts'][0]['Destination'],
            appium_ip=container.attrs['NetworkSettings']['IPAddress'],
            super_host_dir=container.attrs['Mounts'][0]['Source'],
            udid='emulator-5554')

    def run_adbd_as_root(self):
        def check_root_func():
            try:
                # cmd = f'docker exec -t {self.container_name} "adb" ' \
                #       f'"shell" "whoami"'
                cmd = f'docker exec -t {self.container_name} "adb" "root"'
                outs = Utils.run_cmd(cmd)[0]
                # if 'root' in outs:
                if 'adbd is already running as root' in outs:
                    return True
                else:
                    Utils.run_cmd(
                        f'docker exec -t {self.container_name} "adb" "root"')
                    return False
            except:
                logger.debug(traceback.format_exc())
                return False

        if not Utils.wait_until(check_root_func, 10):
            raise DeviceError('failed to run adbd as root')

    def pull_pcap(self, sample_md5: str):
        pcap_path = f'{self.data_dir_on_android}/{sample_md5}.pcap'
        pcap_data = self.webdriver.pull_file(pcap_path)
        with open(f'{self.super_host_dir}/{sample_md5}.pcap', 'wb') as f:
            f.write(base64.b64decode(pcap_data))
        logger.info('pulled pcap successfully.')

    def pull_strace(self, sample_md5: str):
        print('开始pull_strace')
        try:
            strace_path = f'/{self.data_dir_on_android}/{sample_md5[:16]}.strace'            
            strace_data = self.webdriver.pull_file(strace_path)
            # 保存文件
            with open(f'{self.super_host_dir}/{sample_md5}.strace', 'wb') as f:
                f.write(base64.b64decode(strace_data))
            cmd = f'zip -r {self.super_host_dir}/{sample_md5}.strace.zip -j {self.super_host_dir}/{sample_md5}.strace'
            Utils.run_cmd(cmd)
            logger.info('pulled strace successfully.')
            print('拉取strace 成功')
        except Exception as e:
            logger.error(f'拉取 strace 文件时出现错误: {e}')
            print('拉取strace 失败')
            logger.debug(traceback.format_exc())


class Devices:
    def __init__(self):
        self.devices: List[Device] = []
        try:
            self.find_existing_devices()
            logger.info(f'found {len(self.devices)} existing devices.')
            # Create new devices if not enough.
            self.create_new_devices()
            # Availability of docker device should be ensured within each
            # thread because renewing and updating are not well implemented
            # here. However, whether other types of devices' availability
            # should be implemented here are not well considered yet.
            # Annotating following codes is only a quick and dirty solution.
            # available_devices: List[Device] = []
            # # Try to make devices available by restarting
            # while not self.empty():
            #     device = self.devices.pop()
            #     if device.ensure_availability():
            #         available_devices.append(device)
            # # Only available devices are added finally.
            # for device in available_devices:
            #     self.devices.append(device)
            # if self.empty():
            #     raise DeviceError('No any available device.')
            # else:
            #     logger.info(f'{len(self.devices)} devices ready.')
        except Exception as e:
            logger.critical(e)
            logger.debug(traceback.format_exc())

    def find_existing_docker_containers(self):
        container_amount = CFG['device_types']['docker']['amount']
        docker_image = CFG['device_types']['docker']['image']
        try:
            docker_client = docker.from_env()
            containers = docker_client.containers.list(
                all=True, filters={'ancestor': docker_image})
        except requests.exceptions.ConnectionError as e:
            logger.error(f"连接 Docker 时出现网络错误: {e}")
            raise DeviceError("无法连接到 Docker 服务，请检查网络设置。")
        except docker.errors.APIError as e:
            logger.error(f"Docker API 错误: {e}")
            raise DeviceError("Docker API 调用出错，请检查 Docker 服务状态和配置。")
        except docker.errors.DockerException as e:
            logger.error(f"Docker 异常: {e}")
            raise DeviceError("处理 Docker 时出现异常，请检查 Docker 环境。")
        except Exception as e:
            logger.error(f"在处理 Docker 容器时发生未知错误: {e}")
            raise

        unavailable_containers = []
        for container in containers:
            try:
                device: DockerDevice = \
                    DockerDevice.extract_device_from_container(container)
                if device.is_available():
                    self.devices.append(device)
                else:
                    unavailable_containers.append(device)
            except Exception as e:
                logger.warning(f'{e} when adding device from {container.name}.')
                continue
            if len(self.devices) >= container_amount:
                return

        for device in unavailable_containers:
            self.devices.append(device)
            logger.warning(f'{device.container_name} is unavailable but added '
                        f'to existing devices, which will be fixed later.')
            if len(self.devices) >= container_amount:
                return


    def find_existing_devices(self):
        self.find_existing_docker_containers()

    def create_new_devices(self):
        """Create and add to devices a new device."""
        # DockerDevice
        num_existing_containers = \
            sum([1 for d in self.devices if type(d) is DockerDevice])
        container_amount = CFG['device_types']['docker']['amount']
        if num_existing_containers < container_amount:
            for i in range(container_amount - num_existing_containers):
                device = DockerDevice.create()
                if device is not None:
                    self.devices.append(device)

    def get_device(self) -> Device:
        device = self.devices.pop()
        return device

    def empty(self):
        if len(self.devices) == 0:
            return True
        else:
            return False
